import os
import time
from collections import deque
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional, Dict, Any
import uvicorn
from datetime import datetime




# # 快速遍历Windows盘的工具
app = FastAPI(title="Directory Traversal API", 
              description="高效遍历Windows系统目录的API服务", 
              version="1.0.0")

class FileInfo(BaseModel):
    path: str
    name: str
    size: int
    is_dir: bool
    modified_time: str
    created_time: str
    extension: Optional[str] = None

class TraversalResponse(BaseModel):
    count: int
    total_size: int
    traversal_time: float
    files: List[FileInfo]
    directory: str

def format_timestampISO(timestamp: float) -> str:
    """将时间戳转换为ISO 8601格式的字符串"""
    return datetime.fromtimestamp(timestamp).isoformat()

def format_timestamp(timestamp: float) -> str:
    dt = datetime.fromtimestamp(timestamp)
    return dt.strftime("%Y-%m-%d %H:%M:%S.%f")


def traverse_files(root_dir):
    """
    使用广度优先搜索（BFS）非递归遍历目录中的所有文件（含子目录）
    返回生成器以节省内存并提高效率
    """
    queue = deque([root_dir])
    while queue:
        current_dir = queue.popleft()
        try:
            with os.scandir(current_dir) as entries:
                for entry in entries:
                    try:
                        stat = entry.stat()
                        is_dir = entry.is_dir(follow_symlinks=False)
                        
                        file_info = {
                            "path": entry.path,
                            "name": entry.name,
                            "size": stat.st_size if not is_dir else 0,
                            "is_dir": is_dir,
                            "modified_time": format_timestamp(stat.st_mtime),
                            "created_time": format_timestamp(stat.st_ctime),
                            "extension": os.path.splitext(entry.name)[1] if not is_dir and '.' in entry.name else None
                        }
                        
                        if is_dir:
                            queue.append(entry.path)
                        
                        yield file_info
                    except (OSError, PermissionError):
                        # 跳过无法访问的文件或目录
                        continue
        except PermissionError:
            # 跳过权限不足的目录
            continue
        except OSError as e:
            # 处理其他可能的OS错误
            print(f"Error accessing {current_dir}: {e}")
            continue

@app.get("/traverse", response_model=TraversalResponse, 
         summary="遍历指定目录", 
         description="递归遍历指定目录下的所有文件和子目录")
async def traverse_directory(path: Optional[str] = None):
    """
    遍历指定目录并返回所有文件信息
    
    - **path**: 要遍历的目录路径，默认为当前工作目录
    """
    if path is None:
        path = os.getcwd()
    
    # 验证路径是否存在
    if not os.path.exists(path):
        raise HTTPException(status_code=404, detail="指定路径不存在")
    
    # 验证路径是否为目录
    if not os.path.isdir(path):
        raise HTTPException(status_code=400, detail="指定路径不是目录")
    
    start_time = time.time()
    files = []
    total_size = 0
    
    # 遍历目录并收集文件信息
    for file_info in traverse_files(path):
        files.append(file_info)
        if not file_info["is_dir"]:
            total_size += file_info["size"]
    
    traversal_time = time.time() - start_time
    
    return {
        "count": len(files),
        "total_size": total_size,
        "traversal_time": traversal_time,
        "files": files,
        "directory": path
    }


def is_after_target_date(modified_time_str: str, target_date: str = "2025-09-12") -> bool:
    """检查修改时间是否在目标日期之后"""
    # 提取日期部分进行比较
    file_date = modified_time_str.split(" ")[0]  # 获取YYYY-MM-DD部分
    return file_date > target_date



def traverse_filesV2(root_dir, filter_date: str = None):
    """
    使用广度优先搜索（BFS）非递归遍历目录中的所有文件（含子目录）
    返回生成器以节省内存并提高效率
    
    filter_date: 筛选条件，只返回修改时间在此日期之后的文件（格式：YYYY-MM-DD）
    """
    queue = deque([root_dir])
    while queue:
        current_dir = queue.popleft()
        try:
            with os.scandir(current_dir) as entries:
                for entry in entries:
                    try:
                        stat = entry.stat()
                        is_dir = entry.is_dir(follow_symlinks=False)
                        
                        modified_time_str = format_timestamp(stat.st_mtime)
                        
                        # 如果有筛选日期，检查是否满足条件
                        if filter_date and not is_after_target_date(modified_time_str, filter_date):
                            # 不满足筛选条件，跳过此文件
                            if is_dir:
                                queue.append(entry.path)  # 但目录仍需遍历其子目录
                            continue
                        
                        file_info = {
                            "path": entry.path,
                            "name": entry.name,
                            "size": stat.st_size if not is_dir else 0,
                            "is_dir": is_dir,
                            "modified_time": modified_time_str,
                            "created_time": format_timestamp(stat.st_ctime),
                            "extension": os.path.splitext(entry.name)[1] if not is_dir and '.' in entry.name else None
                        }
                        
                        if is_dir:
                            queue.append(entry.path)
                        
                        yield file_info
                    except (OSError, PermissionError):
                        # 跳过无法访问的文件或目录
                        continue
        except PermissionError:
            # 跳过权限不足的目录
            continue
        except OSError as e:
            # 处理其他可能的OS错误
            print(f"Error accessing {current_dir}: {e}")
            continue

@app.get("/traverseFilterDate", response_model=TraversalResponse, 
         summary="遍历指定目录", 
         description="递归遍历指定目录下的所有文件和子目录")
async def traverseFilterDate(path: Optional[str] = None, modified_after: Optional[str] = None):
    """
    遍历指定目录并返回所有文件信息
    
    - **path**: 要遍历的目录路径，默认为当前工作目录
    - **modified_after**: 只返回修改时间在此日期之后的文件（格式：YYYY-MM-DD）
    """
    if path is None:
        path = os.getcwd()
    
    # 验证路径是否存在
    if not os.path.exists(path):
        raise HTTPException(status_code=404, detail="指定路径不存在")
    
    # 验证路径是否为目录
    if not os.path.isdir(path):
        raise HTTPException(status_code=400, detail="指定路径不是目录")
    
    # 验证日期格式
    if modified_after:
        try:
            datetime.strptime(modified_after, "%Y-%m-%d")
        except ValueError:
            raise HTTPException(status_code=400, detail="日期格式不正确，请使用YYYY-MM-DD格式")
    
    start_time = time.time()
    files = []
    total_size = 0
    
    # 遍历目录并收集文件信息
    for file_info in traverse_filesV2(path, modified_after):
        files.append(file_info)
        if not file_info["is_dir"]:
            total_size += file_info["size"]
    
    traversal_time = time.time() - start_time
    
    return {
        "count": len(files),
        "total_size": total_size,
        "traversal_time": traversal_time,
        "files": files,
        "directory": path
    }

@app.get("/stats", summary="获取目录统计信息")
async def get_directory_stats(path: Optional[str] = None, modified_after: Optional[str] = None):
    """
    获取目录的统计信息而不返回所有文件列表
    
    - **path**: 要统计的目录路径，默认为当前工作目录
    - **modified_after**: 只统计修改时间在此日期之后的文件（格式：YYYY-MM-DD）
    """
    if path is None:
        path = os.getcwd()
    
    # 验证路径是否存在
    if not os.path.exists(path):
        raise HTTPException(status_code=404, detail="指定路径不存在")
    
    # 验证路径是否为目录
    if not os.path.isdir(path):
        raise HTTPException(status_code=400, detail="指定路径不是目录")
    
    # 验证日期格式
    if modified_after:
        try:
            datetime.strptime(modified_after, "%Y-%m-%d")
        except ValueError:
            raise HTTPException(status_code=400, detail="日期格式不正确，请使用YYYY-MM-DD格式")
    
    start_time = time.time()
    file_count = 0
    dir_count = 0
    total_size = 0
    
    # 使用os.walk进行更高效的统计（不收集所有文件信息）
    for root, dirs, files_list in os.walk(path):
        # 筛选目录
        if modified_after:
            filtered_dirs = []
            for d in dirs:
                dir_path = os.path.join(root, d)
                try:
                    mtime = os.path.getmtime(dir_path)
                    if datetime.fromtimestamp(mtime).strftime("%Y-%m-%d") > modified_after:
                        filtered_dirs.append(d)
                except (OSError, PermissionError):
                    continue
            dir_count += len(filtered_dirs)
        else:
            dir_count += len(dirs)
        
        # 筛选文件
        if modified_after:
            filtered_files = []
            for file in files_list:
                file_path = os.path.join(root, file)
                try:
                    mtime = os.path.getmtime(file_path)
                    if datetime.fromtimestamp(mtime).strftime("%Y-%m-%d") > modified_after:
                        filtered_files.append(file)
                        total_size += os.path.getsize(file_path)
                except (OSError, PermissionError):
                    continue
            file_count += len(filtered_files)
        else:
            file_count += len(files_list)
            for file in files_list:
                try:
                    file_path = os.path.join(root, file)
                    total_size += os.path.getsize(file_path)
                except (OSError, PermissionError):
                    continue
    
    traversal_time = time.time() - start_time
    
    return {
        "directory": path,
        "file_count": file_count,
        "dir_count": dir_count,
        "total_count": file_count + dir_count,
        "total_size": total_size,
        "traversal_time": traversal_time,
        "filter_date": modified_after
    }




@app.get("/stats", summary="获取目录统计信息")
async def get_directory_stats(path: Optional[str] = None):
    """
    获取目录的统计信息而不返回所有文件列表
    
    - **path**: 要统计的目录路径，默认为当前工作目录
    """
    if path is None:
        path = os.getcwd()
    
    # 验证路径是否存在
    if not os.path.exists(path):
        raise HTTPException(status_code=404, detail="指定路径不存在")
    
    # 验证路径是否为目录
    if not os.path.isdir(path):
        raise HTTPException(status_code=400, detail="指定路径不是目录")
    
    start_time = time.time()
    file_count = 0
    dir_count = 0
    total_size = 0
    
    # 使用os.walk进行更高效的统计（不收集所有文件信息）
    for root, dirs, files in os.walk(path):
        dir_count += len(dirs)
        file_count += len(files)
        for file in files:
            try:
                file_path = os.path.join(root, file)
                total_size += os.path.getsize(file_path)
            except (OSError, PermissionError):
                continue
    
    traversal_time = time.time() - start_time
    
    return {
        "directory": path,
        "file_count": file_count,
        "dir_count": dir_count,
        "total_count": file_count + dir_count,
        "total_size": total_size,
        "traversal_time": traversal_time
    }

@app.get("/", summary="API信息")
async def root():
    return {
        "message": "目录遍历API服务",
        "version": "1.0.0",
        "endpoints": {
            "/traverse": "遍历指定目录下的所有文件和子目录（返回详细信息）",
            "/stats": "获取目录统计信息（不返回文件列表）",
            "/docs": "API文档"
        }
    }

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)